from pydantic import BaseModel, Field
from typing import List, Optional, Union, Dict, Any
from enum import Enum
import time
import uuid


class ModelConfig(BaseModel):
    """Configuration for a model."""
    name: str
    path: Optional[str] = None
    backend: str  # vllm, sglang, transformers
    dtype: str = "auto"
    trust_remote_code: bool = False
    max_tokens: int = 4096
    temperature: float = 1.0
    top_p: float = 1.0


class MessageRole(str, Enum):
    USER = "user"
    ASSISTANT = "assistant"
    SYSTEM = "system"


class ChatMessage(BaseModel):
    """A message in a chat conversation."""
    role: MessageRole
    content: Union[str, List[Dict[str, Any]]]  # Support multimodal content


class Choice(BaseModel):
    """A single choice in a completion response."""
    index: int
    message: ChatMessage
    finish_reason: Optional[str] = "stop"


class Usage(BaseModel):
    """Usage statistics for a completion."""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0


class ChatCompletionRequest(BaseModel):
    """Request schema for chat completion."""
    model: str
    messages: List[ChatMessage]
    temperature: Optional[float] = 1.0
    top_p: Optional[float] = 1.0
    n: Optional[int] = 1
    max_tokens: Optional[int] = 4096
    stop: Optional[Union[str, List[str]]] = None
    stream: Optional[bool] = False
    presence_penalty: Optional[float] = 0
    frequency_penalty: Optional[float] = 0
    logit_bias: Optional[Dict[str, float]] = None


class ChatCompletionResponse(BaseModel):
    """Response schema for chat completion."""
    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex}")
    object: str = "chat.completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str
    choices: List[Choice]
    usage: Usage


class CompletionRequest(BaseModel):
    """Request schema for text completion."""
    model: str
    prompt: Union[str, List[str]]
    suffix: Optional[str] = None
    max_tokens: Optional[int] = 16
    temperature: Optional[float] = 1.0
    top_p: Optional[float] = 1.0
    n: Optional[int] = 1
    stream: Optional[bool] = False
    logprobs: Optional[int] = None
    echo: Optional[bool] = False
    stop: Optional[Union[str, List[str]]] = None
    presence_penalty: Optional[float] = 0
    frequency_penalty: Optional[float] = 0
    best_of: Optional[int] = 1
    logit_bias: Optional[Dict[str, float]] = None
    user: Optional[str] = None


class CompletionResponse(BaseModel):
    """Response schema for text completion."""
    id: str = Field(default_factory=lambda: f"cmpl-{uuid.uuid4().hex}")
    object: str = "text_completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str
    choices: List[Dict[str, Any]]  # Different structure from chat
    usage: Usage